
import argparse
from pathlib import Path
import sys
import numpy as np
from PIL import Image

ALLOWED_IMG_EXT = {".png", ".jpg", ".jpeg", ".tif", ".tiff", ".bmp"}

def parse_color_list(s):
    if not s:
        return []
    out = []
    for part in str(s).split(";"):
        nums = [int(x) for x in part.strip().split(",")]
        if len(nums) != 3:
            raise ValueError(f"Bad color '{part}'. Expected 'R,G,B'.")
        out.append(tuple(nums))
    return out

def to_binary_mask(im: Image.Image, mode="auto", pos_colors=None, gray_thr=128, palette_index=None, strict=False):
    if pos_colors is None: pos_colors = []

    def rgb_to_mask(rgb):
        arr = np.array(rgb)
        if arr.ndim == 2:
            return (arr >= gray_thr).astype(np.uint8)
        if arr.shape[2] == 4:
            arr = arr[:, :, :3]
        H, W, C = arr.shape
        if pos_colors:
            mask = np.zeros((H, W), dtype=np.uint8)
            for col in pos_colors:
                match = (arr[:, :, 0] == col[0]) & (arr[:, :, 1] == col[1]) & (arr[:, :, 2] == col[2])
                mask |= match.astype(np.uint8)
            return mask
        if np.allclose(arr[:, :, 0], arr[:, :, 1]) and np.allclose(arr[:, :, 1], arr[:, :, 2]):
            gray = arr[:, :, 0]
            return (gray >= gray_thr).astype(np.uint8)
        uniq = np.unique(arr.reshape(-1, 3), axis=0)
        if uniq.shape[0] <= 4:
            if any((u == [255, 255, 255]).all() for u in uniq) and any((u == [0, 0, 0]).all() for u in uniq):
                return ((arr[:, :, 0] + arr[:, :, 1] + arr[:, :, 2]) >= 255*3//2).astype(np.uint8)
            if any((u == [255, 0, 0]).all() for u in uniq):
                return ((arr[:, :, 0] == 255) & (arr[:, :, 1] == 0) & (arr[:, :, 2] == 0)).astype(np.uint8)
        if strict:
            raise ValueError("RGB mask ambiguous. Please provide --pos-color.")
        R, G, B = arr[:, :, 0].astype(np.int32), arr[:, :, 1].astype(np.int32), arr[:, :, 2].astype(np.int32)
        return ((R > 128) & (R - np.maximum(G, B) > 40)).astype(np.uint8)

    if mode == "auto":
        if im.mode in ("1", "L"):
            arr = np.array(im.convert("L"))
            return (arr >= gray_thr).astype(np.uint8)
        elif im.mode == "P":
            idx = np.array(im)
            if palette_index is not None:
                return (idx == int(palette_index)).astype(np.uint8)
            if np.unique(idx).size <= 2:
                return (idx != 0).astype(np.uint8)
            if strict:
                raise ValueError("Palette image with >2 indices. Provide --palette-index.")
            vals, counts = np.unique(idx, return_counts=True)
            vals = list(vals); counts = list(counts)
            if 0 in vals:
                zi = vals.index(0); vals.pop(zi); counts.pop(zi)
            if not vals:
                return np.zeros_like(idx, dtype=np.uint8)
            pos = vals[int(np.argmax(counts))]
            return (idx == pos).astype(np.uint8)
        else:
            return rgb_to_mask(im.convert("RGBA"))
    elif mode == "gray":
        arr = np.array(im.convert("L"))
        return (arr >= gray_thr).astype(np.uint8)
    elif mode == "palette":
        idx = np.array(im.convert("P"))
        if palette_index is None:
            raise ValueError("--palette-index is required for mode=palette")
        return (idx == int(palette_index)).astype(np.uint8)
    elif mode == "rgb":
        return rgb_to_mask(im.convert("RGBA"))
    else:
        raise ValueError(f"Unknown mode: {mode}")

def process_split(dataset_dir: Path, split: str, label_dir_name="label",
                  out_label_dir_name="label_bin",
                  mode="auto", pos_colors=None,
                  gray_thr=128, palette_index=None, strict=False, keep_255=True):
    split_dir = dataset_dir / split
    in_label_dir = split_dir / label_dir_name
    out_label_dir = split_dir / out_label_dir_name
    if not in_label_dir.exists():
        raise FileNotFoundError(f"Label dir not found: {in_label_dir}")
    out_label_dir.mkdir(parents=True, exist_ok=True)

    converted = 0
    for p in in_label_dir.rglob("*"):
        if not p.is_file() or p.suffix.lower() not in ALLOWED_IMG_EXT:
            continue
        rel = p.relative_to(in_label_dir)
        out_path = out_label_dir / rel.with_suffix(".png")
        out_path.parent.mkdir(parents=True, exist_ok=True)

        im = Image.open(p)
        try:
            mask01 = to_binary_mask(im, mode=mode, pos_colors=pos_colors,
                                    gray_thr=gray_thr, palette_index=palette_index, strict=strict)
        except Exception as e:
            print(f"[WARN] failed on {p}: {e}", file=sys.stderr)
            continue

        mask = (mask01 * (255 if keep_255 else 1)).astype(np.uint8)
        Image.fromarray(mask, mode="L").save(out_path)
        converted += 1

    return converted, out_label_dir

def maybe_update_lists(dataset_dir: Path, out_label_dir_name="label_bin"):
    lists_dir = dataset_dir / "lists"
    if not lists_dir.exists():
        return
    for list_file in lists_dir.glob("*.txt"):
        lines = []
        changed = False
        for line in list_file.read_text(encoding="utf-8").splitlines():
            if not line.strip():
                continue
            a, b = line.split()
            if "/label/" in b:
                b2 = b.replace("/label/", f"/{out_label_dir_name}/")
                if b2 != b:
                    changed = True
                    b = b2
            elif "\\label\\" in b:
                b2 = b.replace("\\label\\", f"\\{out_label_dir_name}\\")
                if b2 != b:
                    changed = True
                    b = b2
            lines.append(f"{a} {b}")
        if changed:
            list_file.write_text("\n".join(lines) + "\n", encoding="utf-8")
            print(f"[lists] updated {list_file.name} -> {out_label_dir_name}")

def main():
    ap = argparse.ArgumentParser(description="Convert heterogeneous label masks to single-channel binary (0/1 or 0/255).")
    ap.add_argument("--dataset-dir", required=True)
    ap.add_argument("--splits", nargs="+", default=["train", "val", "test"])
    ap.add_argument("--label-dir-name", default="label")
    ap.add_argument("--out-label-dir-name", default="label_bin")
    ap.add_argument("--mode", default="auto", choices=["auto", "gray", "rgb", "palette"])
    ap.add_argument("--pos-color", default=None)
    ap.add_argument("--gray-thr", type=int, default=128)
    ap.add_argument("--palette-index", type=int, default=None)
    ap.add_argument("--strict", action="store_true")
    ap.add_argument("--keep-255", action="store_true")
    ap.add_argument("--update-lists", action="store_true")
    args = ap.parse_args()

    dataset_dir = Path(args.dataset_dir)
    pos_colors = parse_color_list(args.pos_color)

    total_conv = 0
    for split in args.splits:
        m, out_dir = process_split(dataset_dir, split,
                                   label_dir_name=args.label_dir_name,
                                   out_label_dir_name=args.out_label_dir_name,
                                   mode=args.mode, pos_colors=pos_colors,
                                   gray_thr=args.gray_thr, palette_index=args.palette_index,
                                   strict=args.strict, keep_255=args.keep_255)
        print(f"[{split}] converted {m} labels -> {out_dir}")
        total_conv += m

    if args.update_lists:
        maybe_update_lists(dataset_dir, out_label_dir_name=args.out_label_dir_name)

    print(f"Done. Converted {total_conv} files.")

if __name__ == "__main__":
    main()
